from src.module.Model import Model
from src.module.Embedding import Embedding
from langchain_chroma import Chroma

from langchain_core.example_selectors import SemanticSimilarityExampleSelector
from langchain_core.prompts import PromptTemplate, FewShotPromptTemplate


class FlewShortAnswer(Model):
    def start(self, example, prompt, question):
        short_template = FewShotPromptTemplate(
            examples=example,
            example_prompt=prompt,
            suffix='问题: {input}'
        )
        print('short template is:', short_template)
        str_prompt = short_template.invoke({ 'input': question }).to_string()
        print(self.model.invoke(str_prompt))


class ExampleSelector:
    def __init__(self):
        self.flew_short = FlewShortAnswer()
        self.embedding = Embedding().embedding
        self.examples = [
            {
                'question': '你的版本是多少',
                'answer': '版本为 7b llm'
            },
            {
                'question': '你的名字是什么',
                'answer': '我的名字叫小新'
            },
            {
                'question': '小皮球是什么颜色的',
                'answer': '小皮球是蓝色的'
            },
            {
                'question': '今天天气如何',
                'answer': '晴天'
            }
        ]

    def start(self):
        # 1. 将 example 进行向量检索，得出前三条匹配的示例
        example_selector = SemanticSimilarityExampleSelector.from_examples(
            ## 示例内容
            examples = self.examples,
            ## embedding 模型
            embeddings = self.embedding,
            ## 向量数据库
            vectorstore_cls = Chroma,
            ## 只取出匹配度的前三条
            k=3
        )
        search_example = example_selector.select_examples({"question": '小皮球'})

        # 2. 将检索出来之后的 example 数组交给 FlewShort 进行提问, prompt 是对 example 进行格式化，告诉ai为题是什么，答案是什么。
        self.flew_short.start(
            search_example,
            PromptTemplate(
                input_variables=['question', 'answer'],
                template="你现在是一名AI助手，请回答用户的问题 {question}, 答案: {answer}"
            ),
            # PromptTemplate.from_template("你现在是一名AI助手，请回答用户的问题 {question}, 答案: {answer}"), # 与上面的效果是一样的
            "小皮球是什么颜色的")

